{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/MouseLand/cellpose/blob/main/notebooks/test_Cellpose-SAM.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Nb90LCrotIx4"
      },
      "source": [
        "# Cellpose-SAM: superhuman generalization for cellular segmentation\n",
        "\n",
        "Marius Pachitariu, Michael Rariden, Carsen Stringer\n",
        "\n",
        "[paper](https://www.biorxiv.org/content/10.1101/2025.04.28.651001v1) | [code](https://github.com/MouseLand/cellpose)\n",
        "\n",
        "This notebook explains processing example 2D and 3D images using the Cellpose package on Google Colab using the GPU."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z0s2fz5hUk75"
      },
      "source": [
        "### Make sure you have GPU access enabled by going to Runtime -> Change Runtime Type -> Hardware accelerator and selecting GPU\n",
        "\n",
        "![image.png]()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_lRDGixTm1Px"
      },
      "source": [
        "### Install Cellpose-SAM"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "hG3LSmJmLylT"
      },
      "outputs": [],
      "source": [
        "!pip install git+https://www.github.com/mouseland/cellpose.git"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JRalUQBTm1Py"
      },
      "source": [
        "Check GPU and instantiate model - will download weights."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "5ydQ-fggSiUm"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "from cellpose import models, core, io, plot\n",
        "from pathlib import Path\n",
        "from tqdm import trange\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "io.logger_setup() # run this to get printing of progress\n",
        "\n",
        "#Check if colab notebook instance has GPU access\n",
        "if core.use_gpu()==False:\n",
        "  raise ImportError(\"No GPU access, change your runtime\")\n",
        "\n",
        "model = models.CellposeModel(gpu=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fY6Vv5I3m1Py"
      },
      "source": [
        "### Download example images"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "-lZP6alpUAfY"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "from cellpose import utils, io\n",
        "\n",
        "# download example 2D images from website\n",
        "url = \"http://www.cellpose.org/static/data/imgs_cyto3.npz\"\n",
        "filename = \"imgs_cyto3.npz\"\n",
        "utils.download_url_to_file(url, filename)\n",
        "\n",
        "# download 3D tiff\n",
        "url = \"http://www.cellpose.org/static/data/rgb_3D.tif\"\n",
        "utils.download_url_to_file(url, \"rgb_3D.tif\")\n",
        "\n",
        "dat = np.load(filename, allow_pickle=True)[\"arr_0\"].item()\n",
        "\n",
        "imgs = dat[\"imgs\"]\n",
        "masks_true = dat[\"masks_true\"]\n",
        "\n",
        "plt.figure(figsize=(8,3))\n",
        "for i, iex in enumerate([9, 16, 21]):\n",
        "    img = imgs[iex].squeeze()\n",
        "    plt.subplot(1,3,1+i)\n",
        "    plt.imshow(img[0], cmap=\"gray\", vmin=0, vmax=1)\n",
        "    plt.axis('off')\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "M-jKt9wsm1Pz"
      },
      "source": [
        "### Run Cellpose-SAM"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "98DA8zm4A__m"
      },
      "outputs": [],
      "source": [
        "masks_pred, flows, styles = model.eval(imgs, niter=1000) # using more iterations for bacteria\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3JRxBPmatrK7"
      },
      "source": [
        "plot results"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "teNSdi1_m1Pz"
      },
      "outputs": [],
      "source": [
        "from cellpose import transforms, plot\n",
        "\n",
        "titles = [\n",
        "        \"Cellpose\", \"Nuclei\", \"Tissuenet\", \"Livecell\", \"YeaZ\",\n",
        "         \"Omnipose\\nphase-contrast\", \"Omnipose\\nfluorescent\",\n",
        "        \"DeepBacs\"\n",
        "    ]\n",
        "\n",
        "plt.figure(figsize=(12,6))\n",
        "ly = 400\n",
        "for iex in range(len(imgs)):\n",
        "    img = imgs[iex].squeeze().copy()\n",
        "    img = np.clip(transforms.normalize_img(img, axis=0), 0, 1) # normalize images across channel axis\n",
        "    ax = plt.subplot(3, 8, (iex%3)*8 + (iex//3) +1)\n",
        "    if img[1].sum()==0:\n",
        "        img = img[0]\n",
        "        ax.imshow(img, cmap=\"gray\")\n",
        "    else:\n",
        "        # make RGB from 2 channel image\n",
        "        img = np.concatenate((np.zeros_like(img)[:1], img), axis=0).transpose(1,2,0)\n",
        "        ax.imshow(img)\n",
        "    ax.set_ylim([0, min(400, img.shape[0])])\n",
        "    ax.set_xlim([0, min(400, img.shape[1])])\n",
        "\n",
        "\n",
        "    # GROUND-TRUTH = PURPLE\n",
        "    # PREDICTED = YELLOW\n",
        "    outlines_gt = utils.outlines_list(masks_true[iex])\n",
        "    outlines_pred = utils.outlines_list(masks_pred[iex])\n",
        "    for o in outlines_gt:\n",
        "        plt.plot(o[:,0], o[:,1], color=[0.7,0.4,1], lw=0.5)\n",
        "    for o in outlines_pred:\n",
        "        plt.plot(o[:,0], o[:,1], color=[1,1,0.3], lw=0.75, ls=\"--\")\n",
        "    plt.axis('off')\n",
        "\n",
        "    if iex%3 == 0:\n",
        "        ax.set_title(titles[iex//3])\n",
        "\n",
        "plt.tight_layout()\n",
        "plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rMyZtY6ym1P0"
      },
      "source": [
        "# Run Cellpose-SAM in 3D\n",
        "\n",
        "There are two ways to run cellpose in 3D, this cell shows both, choose which one works best for you.\n",
        "\n",
        "First way: computes flows from 2D slices and combines into 3D flows to create masks\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2b2hVxCvm1P0"
      },
      "outputs": [],
      "source": [
        "img_3D = io.imread(\"rgb_3D.tif\")\n",
        "\n",
        "\n",
        "# 1. computes flows from 2D slices and combines into 3D flows to create masks\n",
        "masks, flows, _ = model.eval(img_3D, z_axis=0, channel_axis=1,\n",
        "                                batch_size=32,\n",
        "                                do_3D=True, flow3D_smooth=1)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KeMtAuRom1P0"
      },
      "source": [
        "Second way: computes masks in 2D slices and stitches masks in 3D based on mask overlap\n",
        "\n",
        "Note stitching (with stitch_threshold > 0) can also be used to track cells over time."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "WTyCgBUfm1P0"
      },
      "outputs": [],
      "source": [
        "# 2. computes masks in 2D slices and stitches masks in 3D based on mask overlap\n",
        "print('running cellpose 2D + stitching masks')\n",
        "masks_stitched, flows_stitched, _ = model.eval(img_3D, z_axis=0, channel_axis=1,\n",
        "                                                  batch_size=32,\n",
        "                                                  do_3D=False, stitch_threshold=0.5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wbu1j0h6m1P0"
      },
      "source": [
        "Results from 3D flows => masks computation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Vfg67u2dm1P0"
      },
      "outputs": [],
      "source": [
        "# DISPLAY RESULTS 3D flows => masks\n",
        "plt.figure(figsize=(15,3))\n",
        "for i,iplane in enumerate(np.arange(0,75,10,int)):\n",
        "  img0 = plot.image_to_rgb(img_3D[iplane, [1,0]].copy(), channels=[2,3])\n",
        "  plt.subplot(1,8,i+1)\n",
        "  outlines = utils.masks_to_outlines(masks[iplane])\n",
        "  outX, outY = np.nonzero(outlines)\n",
        "  imgout= img0.copy()\n",
        "  imgout[outX, outY] = np.array([255,75,75])\n",
        "  plt.imshow(imgout)\n",
        "  plt.title('iplane = %d'%iplane)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dj18ZyzHm1P0"
      },
      "source": [
        "Results from stitching"
      ]
    },
    {
      "cell_type": "code",
      "source": [
        "# DISPLAY RESULTS stitching\n",
        "plt.figure(figsize=(15,3))\n",
        "for i,iplane in enumerate(np.arange(0,75,10,int)):\n",
        "  img0 = plot.image_to_rgb(img_3D[iplane, [1,0]].copy(), channels=[2,3])\n",
        "  plt.subplot(1,8,i+1)\n",
        "  outlines = utils.masks_to_outlines(masks_stitched[iplane])\n",
        "  outX, outY = np.nonzero(outlines)\n",
        "  imgout= img0.copy()\n",
        "  imgout[outX, outY] = np.array([255,75,75])\n",
        "  plt.imshow(imgout)\n",
        "  plt.title('iplane = %d'%iplane)"
      ],
      "metadata": {
        "id": "fd-6Hji-n9_H"
      },
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "provenance": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "display_name": "cellpose",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.11.9"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}